import numpy as np
import matplotlib.pyplot as plt

# 定义 SH 曲线函数
def SH1_CURVE(t, x4):
    if t <= 0.0:
        return 0.0
    elif t < x4:
        return (t / x4) ** 2.5
    else:
        return 1.0

def SH2_CURVE(t, x4):
    if t <= 0.0:
        return 0.0
    elif t <= x4:
        return 0.5 * (t / x4) ** 2.5
    elif t < 2 * x4:
        return 1 - 0.5 * (2 - t / x4) ** 2.5
    else:
        return 1.0

# 读取参数和数据
other_para = np.loadtxt('others.txt')
data = np.loadtxt('inputData.txt')

area = other_para[0]
upperTankRatio = other_para[1]
lowerTankRatio = other_para[2]

P = data[:, 0]
E = data[:, 1]
Qobs = data[:, 2]
Qobs_mm = Qobs * 86.4 / area  # m³/s -> mm/d
nStep = data.shape[0]

# 使用深度学习优化得到的最佳参数组合
x1 = 108.8
x2 = -2.30
x3 = 98.2
x4 = 1.07

print(f"Running GR4J model with optimal parameters:")
print(f"x1 = {x1}")
print(f"x2 = {x2}")
print(f"x3 = {x3}")
print(f"x4 = {x4}")
print(f"Expected NSE: 0.860")

# 初始化
Pn = np.zeros(nStep)
En = np.zeros(nStep)
Ps = np.zeros(nStep)
Es = np.zeros(nStep)
Perc = np.zeros(nStep)
Pr = np.zeros(nStep)
Qr = np.zeros(nStep)
Qd = np.zeros(nStep)
Q = np.zeros(nStep)

maxDayDelay = 10
SH1 = np.array([SH1_CURVE(i + 1, x4) for i in range(maxDayDelay)])
SH2 = np.array([SH2_CURVE(i + 1, x4) for i in range(2 * maxDayDelay)])
UH1 = np.diff(np.concatenate(([0], SH1)))
UH2 = np.diff(np.concatenate(([0], SH2)))

for i in range(nStep):
    if P[i] >= E[i]:
        Pn[i] = P[i] - E[i]
        En[i] = 0
    else:
        Pn[i] = 0
        En[i] = E[i] - P[i]

S0 = upperTankRatio * x1
R0 = lowerTankRatio * x3
S = np.zeros(nStep)
R = np.zeros(nStep)
UH_Fast = np.zeros((nStep, maxDayDelay))
UH_Slow = np.zeros((nStep, 2 * maxDayDelay))
S_TEMP = S0
R_TEMP = R0

for i in range(nStep):
    S[i] = S_TEMP
    R[i] = R_TEMP

    if Pn[i] != 0:
        Ps[i] = x1 * (1 - (S[i] / x1) ** 2) * np.tanh(Pn[i] / x1) / (1 + S[i] / x1 * np.tanh(Pn[i] / x1))
        Es[i] = 0
    elif En[i] != 0:
        Ps[i] = 0
        Es[i] = (S[i] * (2 - (S[i] / x1)) * np.tanh(En[i] / x1)) / (1 + (1 - S[i] / x1) * np.tanh(En[i] / x1))

    S_TEMP = S[i] - Es[i] + Ps[i]
    Perc[i] = S_TEMP * (1 - (1 + (4.0 / 9.0 * (S_TEMP / x1)) ** 4) ** (-0.25))
    Pr[i] = Perc[i] + (Pn[i] - Ps[i])
    S_TEMP -= Perc[i]

    F = x2 * (R[i] / x3) ** 3.5
    R_Fast = Pr[i] * 0.9
    R_Slow = Pr[i] * 0.1

    UH_Fast[i, :] = R_Fast * UH1
    UH_Slow[i, :] = R_Slow * UH2

    if i > 0:
        UH_Fast[i, :-1] += UH_Fast[i - 1, 1:]
        UH_Slow[i, :-1] += UH_Slow[i - 1, 1:]

    R_TEMP = max(0, R_TEMP + UH_Fast[i, 0] + F)
    Qr[i] = R_TEMP * (1 - (1 + (R_TEMP / x3) ** 4) ** (-0.25))
    R_TEMP -= Qr[i]
    Qd[i] = max(0, UH_Slow[i, 0] + F)
    Q[i] = Qr[i] + Qd[i]

# 计算NSE值进行验证（跳过前365天）
Qobs_eval = Qobs_mm[365:]
Q_eval = Q[365:]
Q_mean = np.mean(Qobs_eval)
numerator = np.sum((Qobs_eval - Q_eval) ** 2)
denominator = np.sum((Qobs_eval - Q_mean) ** 2)
NSE = 1 - numerator / denominator if denominator != 0 else -np.inf

print(f"Calculated NSE: {NSE:.4f}")

# Plot results
plt.figure(figsize=(14, 8))

# Main plot - Complete time series
plt.subplot(2, 1, 1)
plt.plot(range(nStep), Qobs_mm, label="Observed", color="black", linewidth=1.5, alpha=0.8)
plt.plot(range(nStep), Q, label="Simulated", color="red", linestyle='-', linewidth=1.2, alpha=0.9)
plt.title(f"GR4J Model Simulation Results - Optimal Parameters\n(x1={x1}, x2={x2}, x3={x3}, x4={x4}, NSE={NSE:.4f})", fontsize=14)
plt.xlabel("Time (days)", fontsize=12)
plt.ylabel("Flow (mm/d)", fontsize=12)
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)

# Subplot - Detailed view (first 2 years)
plt.subplot(2, 1, 2)
days_to_show = min(730, nStep)  # Show first 2 years or all data
plt.plot(range(days_to_show), Qobs_mm[:days_to_show], label="Observed", color="black", linewidth=1.5, alpha=0.8)
plt.plot(range(days_to_show), Q[:days_to_show], label="Simulated", color="red", linestyle='-', linewidth=1.2, alpha=0.9)
plt.title("Detailed Comparison - First Two Years", fontsize=12)
plt.xlabel("Time (days)", fontsize=12)
plt.ylabel("Flow (mm/d)", fontsize=12)
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# Print model performance statistics
print("\n=== Model Performance Statistics ===")
print(f"Total simulation steps: {nStep}")
print(f"Evaluation steps (skip first 365 days): {len(Qobs_eval)}")
print(f"Observed flow mean: {np.mean(Qobs_mm):.3f} mm/d")
print(f"Simulated flow mean: {np.mean(Q):.3f} mm/d")
print(f"Observed flow std: {np.std(Qobs_mm):.3f} mm/d")
print(f"Simulated flow std: {np.std(Q):.3f} mm/d")
print(f"Nash-Sutcliffe Efficiency (NSE): {NSE:.4f}")

# Calculate other evaluation metrics
correlation = np.corrcoef(Qobs_eval, Q_eval)[0, 1]
rmse = np.sqrt(np.mean((Qobs_eval - Q_eval) ** 2))
mae = np.mean(np.abs(Qobs_eval - Q_eval))

print(f"Correlation coefficient (R): {correlation:.4f}")
print(f"Root Mean Square Error (RMSE): {rmse:.3f} mm/d")
print(f"Mean Absolute Error (MAE): {mae:.3f} mm/d") 